# p is generated as 2-dim
import csv
import datetime
import os
from tabnanny import filename_only

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MaxNLocator
from openpyxl import Workbook
from sko.GA import GA
from sko.tools import set_run_mode

# to monitor optimization process
from osc02.optimization.get_cost import get_cost_with_simulation, get_sampled_duration_and_delivery_time
from osc02.static.constants import Constants, Debugs, DirectoryConfig

iter_counter = 1


def fitness_func(start_correction):
    # print("iterating ...", Constants.N_ACTIVITY, start_correction)
    # print(type(start_correction))
    return get_cost_with_simulation(start_correction=start_correction,
                                    seed=seed,
                                    unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
                                    unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
                                    unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
                                    project_deadline=Debugs.project_deadline
                                    )


# %%
# 保证目录存在
DirectoryConfig.ensure_all_dirs_exist()
filename_summary = os.path.join(DirectoryConfig.OUTPUT_DIR, 'summary.xlsx')
# summary_filename = 'summary.xlsx'
if os.path.exists(filename_summary):
    os.remove(filename_summary)

wb_summary = Workbook()
# ws = wb.remove_sheet
ws_summary = wb_summary.active

# prob_mut_dict = {1: 1e-05, 4: 6e-05}
# to try different prob_mut. But, prob_mut is set static. The loop is useless.
# seed is NOT used in this program but in subs.
# it just is a STUD
seed = 1
# for unit_component_late_acceptance_penalty_by_day in [500, 1000, 2000]:
#     for unit_project_late_completion_penalty_by_day in [2000, 5000, 10000]:
filename_sensitivity_table = os.path.join(DirectoryConfig.OUTPUT_DIR,"sensitivity_table.csv")


def visualize():
    y_history = pd.DataFrame(ga.all_history_Y)
    plt.rcParams['font.family'] = 'Arial'
    plt.figure(figsize=(10, 8), dpi=550)
    # plt.subplots_adjust(bottom=0.2)  # 增加底部空间
    # plt.tight_layout()
    fig, ax = plt.subplots(2, 1)
    ax[0].plot(y_history.index, y_history.values, '.', color='blue')
    ax[0].set_xlabel('Generation')
    ax[0].set_ylabel('Penalty value')
    ax[0].set_title("Genetic algorithm iteration\n" + 'pop=' + str(Debugs.size_pop)
                    + ', max_iter=' + str(Debugs.max_iter) + ', best=' + str(round(best_y[0], 2))
                    # + ',prob_mut=' + str(round(prob_mut, 10))
                    )
    # Y_history.min(axis=1).cummin().plot(kind='line')
    ax[1].bar(range(1, Constants.N_ACTIVITY + 1), np.array(best_x), color='blue')
    ax[1].set_ylim(0, 7)
    ax[1].set_xlabel(f'Activities\' no\n\n'
                     f'CLAP={unit_component_late_acceptance_penalty_by_day}'
                     f'  PLCP={unit_project_late_completion_penalty_by_day}')
    ax[1].set_ylabel('Postponement / day')
    ax[1].set_title("Suggested postponements on the baseline schedule")
    # plt.legend()
    # make locator integer
    plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))
    # fig.savefig('line plot.jpg', bbox_inches='tight', dpi=150)
    # filename_convergence_postponement = 'Optimization_CLAP' + str(unit_component_late_acceptance_penalty_by_day).zfill(4) + '_PLCP' + str(
    #     unit_project_late_completion_penalty_by_day).zfill(5) + '.png'
    # filename_convergence_postponement = os.path.join(DirectoryConfig.OUTPUT_DIR,filename_convergence_postponement)
    filename_convergence_postponement = os.path.join(
        DirectoryConfig.OUTPUT_DIR,
        f'Optimization_CLAP{unit_component_late_acceptance_penalty_by_day:04d}_PLCP{unit_project_late_completion_penalty_by_day:05d}.png'
    )
    fig.savefig(filename_convergence_postponement,
                # bbox_inches='tight',
                dpi=150)
    plt.show()


def getPostponementCost(x):
    # global total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty, total_cost
    avg_component_late_acceptance_penalty, avg_component_late_delivery_penalty, project_late_completion_penalty \
        = get_cost_with_simulation(start_correction=x,
                                   seed=seed,
                                   unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
                                   unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
                                   unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
                                   project_deadline=Debugs.project_deadline,
                                   debug=True
                                   )
    total_cost = avg_component_late_acceptance_penalty \
                 - avg_component_late_delivery_penalty \
                 + project_late_completion_penalty
    return avg_component_late_acceptance_penalty, avg_component_late_delivery_penalty, project_late_completion_penalty, total_cost


# When debugging the format, use a minimal parameter to reduce execution time
# Debugs.max_iter = 1
# Debugs.size_pop = 2
#
with open(filename_sensitivity_table, 'a', newline='') as sensitivity_table_csv_file:
    sensitivity_writer = csv.writer(sensitivity_table_csv_file)
    header = [
        'Unit CLAP/day [$]',
        'Unit PLCP/day [$]',
        'Ratio of unit PLCP to unit CLAP',
        'Avg. TAC of B.S. [$]',
        'Avg. CLAP of B.S. [$]',
        'Avg. CLDP of B.S. [$]',
        'Avg. PLCP of B.S. [$]',
        # 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
        'Avg. suggested postponement [day]',
        'Avg. reduced TAC [$]',
        'Avg. reduced CLAP [$]',
        'Avg. reduced CLDP [$]',
        'Avg. reduced PLCP [$]',
        'Reduction rate'
    ]
    sensitivity_writer.writerow(header)
    # for unit_component_late_acceptance_penalty_by_day in [500, 1000, 2000]:
    for unit_component_late_acceptance_penalty_by_day in [1000]:
        # for unit_component_late_acceptance_penalty_by_day in [500]:
        # for unit_project_late_completion_penalty_by_day in [2000, 5000, 10000]:
        for unit_project_late_completion_penalty_by_day in [5000]:
            # for unit_project_late_completion_penalty_by_day in [2000]:
            Debugs.unit_component_late_acceptance_penalty_by_day = \
                unit_component_late_acceptance_penalty_by_day
            Debugs.unit_project_late_completion_penalty_by_day = \
                unit_project_late_completion_penalty_by_day
            # prob_mut = prob_mut_dict[seed]
            prob_mut = 1e-06
            # if seed == 6:
            #     prob_mut = 5e-04
            filename_receiving_plan = 'Receiving plan under different waiting penalty level' \
                                      + str(seed) + '-' + str(round(prob_mut, 10)) + '.xlsx'
            filename_receiving_plan = os.path.join(DirectoryConfig.OUTPUT_DIR, filename_receiving_plan)
            if os.path.exists(filename_receiving_plan):
                os.remove(filename_receiving_plan)

            wb = Workbook()
            # ws = wb.remove_sheet
            ws = wb.active
            # ws = wb.create_sheet('sheet')
            lb_element = 0
            ub_element = 7
            ws.append(['size_pop=', Debugs.size_pop, 'max_iter=', Debugs.max_iter, 'ub=', ub_element])
            title = [' ', 'seed', 'unit_waiting_penalty', 'unit_project_late_completion_penalty', 'total cost']

            for i in range(1, Constants.N_ACTIVITY + 1):
                title.append('C' + str(i))
            ws.append(title)

            # set_run_mode(get_cost, 'vectorization')
            # set_run_mode(get_cost, 'parallel')
            set_run_mode(fitness_func, 'cached')

            ga = GA(func=fitness_func, n_dim=Constants.N_ACTIVITY,
                    size_pop=Debugs.size_pop,
                    max_iter=Debugs.max_iter,
                    prob_mut=prob_mut,  # probability of mutation
                    lb=[lb_element] * Constants.N_ACTIVITY,
                    ub=[ub_element] * Constants.N_ACTIVITY,
                    # precision=[1] * Constants.N_ACTIVITY)
                    precision=1
                    )

            # ga.Chrom = np.random.randint(0, 1, size=(size_pop, Constants.N_ACTIVITY))
            # manually set the initial population, preset [0 0 0.. 0 0] * 18 * 3 as a chromosome
            inital_pop = np.random.randint(0, 7, size=(Debugs.size_pop, Constants.N_ACTIVITY * 3))
            # set one of individual as [0] * Constants.N_ACTIVITY * 3
            inital_pop[0] = [0] * Constants.N_ACTIVITY * 3
            ga.Chrom = inital_pop

            act_delivery_delay, durations = get_sampled_duration_and_delivery_time(seed=seed)
            print( 'act_delivery_delay ', act_delivery_delay)
            print('duration ', durations)
            base_total_waiting_cost = get_cost_with_simulation(start_correction=[0] * Constants.N_ACTIVITY,
                                                               seed=seed,  # NOT used in the invoked function
                                                               unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
                                                               unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
                                                               unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
                                                               project_deadline=Debugs.project_deadline
                                                               )
            # Sampled data are changed to created in 20 scenarios, seems durations  are not suitable here
            row = ['base cost', seed, Debugs.unit_component_late_acceptance_penalty_by_day,
                   Debugs.unit_project_late_completion_penalty_by_day]
            row.extend([base_total_waiting_cost])
            row.extend(durations)
            ws_summary.append(row)

            ws.append(row)
            row = ['delivery delay', seed, Debugs.unit_component_late_acceptance_penalty_by_day,
                   Debugs.unit_project_late_completion_penalty_by_day]
            row.extend(' ')
            row.extend(act_delivery_delay)
            ws.append(row)

            start_time = datetime.datetime.now()
            # run ga
            best_x, best_y = ga.run()
            # print("ga.run()", iter_counter)
            iter_counter += 1
            # print('best_x:', best_x, '\n', 'best_y:', best_y)
            lapse_time = (datetime.datetime.now() - start_time).total_seconds()
            ws.append(['lapse', lapse_time])
            row = ['min cost', seed, Debugs.unit_component_late_acceptance_penalty_by_day,
                   Debugs.unit_project_late_completion_penalty_by_day]
            print('best_x=', best_x)
            print('best_y=', best_y)
            row.extend(best_y)
            row.extend(best_x)
            ws.append(row)
            ws_summary.append(row)
            # After solving the problem,plug the solution, best_x to plot
            # Keep temporary data for debug
            print("After solving the problem,plug the solution, best_x to plot")
            # get all the costs of best_x
            avg_component_late_acceptance_penalty, \
            avg_component_late_delivery_penalty, \
            avg_project_late_completion_penalty, \
            total_cost = getPostponementCost(best_x)

            details = ['total cost = ', total_cost,
                       'total_waiting_cost=', round(avg_component_late_acceptance_penalty, 2),
                       'total_late_delivery_penalty=', round(avg_component_late_delivery_penalty, 2),
                       'project_late_completion_penalty=', round(avg_project_late_completion_penalty, 2)]
            ws.append(details)
            ws.append([])
            ws_summary.append(details)
            wb.save(filename_receiving_plan)
            wb.close()
            print('seed=', seed, 'lapse_time=', lapse_time)

            visualize()

            # get all the costs of base schedule ([0] * 18)
            avg_component_late_acceptance_penalty_on_base_schedule, \
            avg_component_late_delivery_penalty_on_base_schedule, \
            avg_project_late_completion_penalty_on_base_schedule, \
            avg_total_additional_cost_of_base_schedule = getPostponementCost([0] * 18)

            combined_row = [
                unit_component_late_acceptance_penalty_by_day,  # unit CLAP
                unit_project_late_completion_penalty_by_day,  # # PLCP
                (unit_project_late_completion_penalty_by_day /
                 unit_component_late_acceptance_penalty_by_day),  # Ratio of unit CLAP to unit PLCP
                avg_total_additional_cost_of_base_schedule,  # Avg. TAC
                avg_component_late_acceptance_penalty_on_base_schedule,  # CLAP
                avg_component_late_delivery_penalty_on_base_schedule,  # CLDP
                avg_project_late_completion_penalty_on_base_schedule,  # PLCP
                round(sum(best_x) / len(best_x), 2),  # avg. postponement
                best_y[0],  # reduced TAC, best_y is a ndarray
                round(avg_component_late_acceptance_penalty, 2),  # CLAP
                round(avg_component_late_delivery_penalty, 2),  # CLDP
                round(avg_project_late_completion_penalty, 2),  # PLCP
                round((avg_total_additional_cost_of_base_schedule - best_y[0])
                      / avg_total_additional_cost_of_base_schedule
                      , 2)
            ]
            sensitivity_writer.writerow(combined_row)
        # prob_mut *= 0.4

wb_summary.save(filename_summary)
wb_summary.close()
print("End of optimization")
# base schedule produces an average of additional cost on scenarios

# x = [0] * Constants.N_ACTIVITY
# total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty \
#     = get_cost_with_simulation(x,
#                                seed=None,
#                                unit_waiting_penalty=Debugs.unit_waiting_penalty,
#                                unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty,
#                                unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty,
#                                project_deadline=Debugs.project_deadline,
#                                debug=True
#                                )
# total_cost = total_waiting_cost - total_late_delivery_penalty + project_late_completion_penalty
# print("x=", x, "total_cost=", total_cost)
